// KSP.cpp : Defines the entry point for the console application.
//

#include <iostream>
#include <time.h>

#include "ProgramOptions.h"
#include "Batch.h"

#include <vbl/img/ImageDraw.h>

struct cutOutInfo 
{ 
	rectangle r;
	int pathNr;
	img::ImageRGB imageCutOut;
};

//typedef ProgramOptions PO;

void getSourceNodes( vector<vector<location>> & complete_paths, int last_time_step, vector<location> & sources )
{
	if ( complete_paths.size() != 0 ) 
	{
		for( unsigned int i = 0 ; i < complete_paths.size() ; i++ ) 
		{
			vector<location> path = complete_paths[i];
			location l = path[path.size() - 2];
			if( l.ts == last_time_step ) 
			{
				sources.push_back(l);
			}
		}
	}
}

location intToLocation( int loc , int numOfLocations , int pomX ) 
{
	int ts = (int) floor( ((double) loc-1) / numOfLocations );
	loc = ( ( loc - 1 ) % numOfLocations ) + 1;
	int locx = ((loc-1) % pomX )+1;
	int locy = (( loc - locx ) / pomX )+1;
	
	location l;
	l.ts = ts;
	l.x = locx;
	l.y = locy;
	return l;
}

void combinePaths( PO & po, vector<vector<int>> & partialPaths, vector<vector<location>> & completePaths, int lastTimeStep ) 
{
	if ( completePaths.size() == 0 ) 
	{
		for ( unsigned int i = 0 ; i < partialPaths.size() ; i++ ) 
		{
			vector<int> partialPath = partialPaths[i];
			vector<location> toAdd;
			for( unsigned int j = 0 ; j < partialPath.size() ; j++ ) 
			{
				location l = intToLocation(partialPath[j] , po.pom_size, po.pom_x );
				toAdd.push_back(l);
			}
			completePaths.push_back(toAdd);
		}
	}
	else 
	{
		for ( unsigned int i = 0 ; i < partialPaths.size() ; i++ ) 
		{
			int extensionFound = -1;
			vector<int> partialPath = partialPaths[i];
			int l = partialPath[1];
			location loc = intToLocation(l , po.pom_size, po.pom_x );
			vector<location> completePath;
			if ( loc.ts == 0 ) 
			{
				// path starts at beginning of batch
				for ( unsigned int j = 0 ; j < completePaths.size() ; j++ ) 
				{
					completePath = completePaths[j];
					location lastLoc = completePath[completePath.size() - 2];
					if ( lastLoc.ts == lastTimeStep && lastLoc.x == loc.x && lastLoc.y == loc.y ) 
					{
						extensionFound = j;
						break;
					}
				}
			}
			if ( extensionFound != -1 ) 
			{
				completePaths.erase( completePaths.begin() + extensionFound );
				location last = completePath[ completePath.size() - 1 ];
				int ts = last.ts;
				completePath.pop_back();
				for ( unsigned int j = 2 ; j < partialPath.size() ; j++ ) 
				{
					location toAdd = intToLocation( partialPath[j], po.pom_size, po.pom_x );
					toAdd.ts = ts;
					ts++;
					completePath.push_back(toAdd);
				}
				completePaths.push_back(completePath);
			} 
			else 
			{
				vector<location> toAdd;
				for( unsigned int j = 0 ; j < partialPath.size() ; j++ ) 
				{
					location l = intToLocation(partialPath[j], po.pom_size, po.pom_x );
					l.ts += lastTimeStep;
					toAdd.push_back(l);
				}
				completePaths.push_back(toAdd);
			}
		}
	}
}

void savePaths( PO & po, vector<vector<location>> & paths , vector<vector<int>> & fullPaths )
{
	
	for( unsigned int i = 0 ; i < paths.size() ; i++ )
	{
		vector<int> path;
		for( unsigned int j = 0 ; j < po.num_of_frames ; j++ ) 
		{
			path.push_back(0);
		}
		fullPaths.push_back(path);
	}

	for( unsigned int i = 0 ; i < paths.size() ; i++ ) 
	{
		vector<location> path = paths[i];
		for( unsigned int j = 1 ; j < path.size()-1 ; j++ ) 
		{
			int loc = (path[j].y-1) * po.pom_x + path[j].x;
			int ts = path[j].ts;
			//int loc = ( ( path[j] - 1 ) % numOfLocations ) + 1;
			//int ts = (int) floor( ((double) path[j]-1) / numOfLocations );
			fullPaths[i][ts] = loc;
		}
	}

	ofstream outFile;
	string fileLoc = po.output_path + "paths.txt";
	outFile.open(fileLoc.c_str());
	for( unsigned int i = 0 ; i < fullPaths.size() ; i++ ) 
	{
		vector<int> path = fullPaths[i];
		for( unsigned int j = 0 ; j < path.size() ; j++ ) 
		{
			outFile << path[j] << " ";
		}
		outFile << endl;
	}
	outFile.close();
}

img::Triple<> getHeatmapColor( double value ) 
{
	int r, g, b;

	if ( value > 0.75 && value <= 1 )
	{
		r = 255;
		b = 0;
		g = int((1-((value-0.75) * 4)) * 255);
	} 
	else if ( value > 0.5 ) 
	{
		g = 255;
		b = 0;
		r = int(((value-0.5)*4) * 255);
	} 
	else if ( value > 0.25 ) 
	{
		r = 0;
		g = 255;
		b = int((1-((value-0.25) * 4)) * 255);
	} 
	else if ( value > 0 ) 
	{
		r = 0;
		b = 255;
		g = int(value*4*255);
	}

	return img::Triple<>( r , g , b );
}

void saveImages( PO po, vector<vector<int>> & paths )
{
	cout << "Writing images to disc." << endl;

	vector< img::Triple<> > colors;
	colors.push_back(img::Triple<>(255,0,0));	// red
	colors.push_back(img::Triple<>(0,255,0));	// green
	//colors.push_back(img::Triple<>(0,0,255));	// blue
	colors.push_back(img::Triple<>(255,255,0)); // yellow
	colors.push_back(img::Triple<>(0,255,255));	// cyan
	colors.push_back(img::Triple<>(255,0,255)); // purple
	colors.push_back(img::Triple<>(255,165,0)); // orange
	colors.push_back(img::Triple<>(165,42,42)); // brown
	colors.push_back(img::Triple<>(0,255,127)); // spring green
	colors.push_back(img::Triple<>(255,192,203)); // pink
	colors.push_back(img::Triple<>(127,255,212)); // aquamarine
	colors.push_back(img::Triple<>(218,165,32)); // goldenrod
	colors.push_back(img::Triple<>(178,34,34)); // firebrick
	colors.push_back(img::Triple<>(205,133,63)); // peru
	colors.push_back(img::Triple<>(238,130,238)); // violet

	vector< img::ImageRGB > inputIms( po.num_of_cams , img::ImageRGB( po.image_width, po.image_height ) );
	
	boost::progress_display progress( po.num_of_frames );
	for ( int i = 0 ; i < po.num_of_frames ; i++ ) 
	{
		vector< img::ImageRGB > imageVec;

		int frame = po.begin_frame + i;
		string frameNumber;
		if ( frame < 1000 ) {
			frameNumber = "0" + lexical_cast<string>(frame);
		} else {
			frameNumber = lexical_cast<string>(frame);
		}

		for( int cam = 0; cam < po.num_of_cams; ++cam ) 
		{
			img::ImageG imRaw( po.image_width, po.image_height );
#ifndef USE_OLD_FILESYSTEM
			img::io::readRAW( po.files_per_camera[cam][ i ], imRaw, imRaw.width(), imRaw.height() ); 
#else
			img::io::readRAW( po.image_path + "/" + po.files_per_camera[cam][ po.begin_frame + i ], imRaw, imRaw.width(), 
				imRaw.height() ); 
#endif
			img::proc::demosaicPixelGrouping( imRaw, inputIms[cam] );
			po.undistortImage( inputIms[cam], po.dist_map.first[cam], po.dist_map.second[cam] );
			img::rotateSafe< img::InterpolatorBilinear >( inputIms[cam], po.image_width/2, po.image_height/2, po.rot[cam] );
			img::ImageRGB im = inputIms[cam];

			img::ImageRGB background;
#ifndef USE_CAS_DATASET
			string backgroundImPath = po.background_path + lexical_cast<string>(frame) + "_" + lexical_cast<string>(cam) + 
				"-result.png";
#else
			string backgroundImPath = po.background_path + "result-f" + lexical_cast<string>(frame) + "-c" + 
				lexical_cast<string>(cam) + ".png";
#endif
			try
			{
				img::io::read( backgroundImPath , background );
			} catch( vbl::runtime_error e ) 
			{
				cout << " Could not read background file: " << e.what() << endl;
				exit(1);
			}
			img::Triple<> maskColor(0,204,153);
			
			for ( unsigned int j = 0 ; j < po.pom_locations[cam].size() ; j++ ) {
				rectangle r = po.pom_locations[cam][j];
				double value = po.probabilities[i][j];
				img::Triple<> color = getHeatmapColor( value );
				if ( r.isVisible ) {
					img::drawCross( im , int((r.xmin+r.xmax) / 2.0) , r.ymax, color, 3, 3 );
				}
			}

			vector<pair<int,int>> toSort;
			for ( unsigned int j = 0 ; j < paths.size() ; j++ ) {
				int loc = paths[j][i];
				if ( loc != 0 ) {
					rectangle r = po.pom_locations[cam][loc-1];
					if ( r.xmax != -1 ) {
						int surface = po.calculateSurface(r);
						toSort.push_back(pair<int,int>(surface,j));
					}
				}
			}
			sort (toSort.begin(), toSort.end());

			vector<cutOutInfo> toAdd;
			for ( unsigned int j = 0 ; j < toSort.size() ; j++ ) {
				int loc = paths[toSort[j].second][i];
				if ( loc != 0 ) {
					rectangle r = po.pom_locations[cam][loc-1];
					img::ImageRGB backgroundCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::ImageRGB imageCutOut( r.xmax - r.xmin , r.ymax - r.ymin );
					img::Triple<> pathColor = colors[ (toSort[j].second % colors.size() )];
					try
					{
						img::copy( background , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , backgroundCutOut );
						img::equals( backgroundCutOut, maskColor, pathColor, img::Triple<>(0,0,0) );
						img::copy( im , r.xmin , r.ymin , r.xmax - r.xmin , r.ymax - r.ymin , imageCutOut);
						overlay(imageCutOut,backgroundCutOut, 0.5, true);	
					} catch ( vbl::runtime_error e )
					{
						cout << "Unable overlay part of background onto part of image. Error: " << e.what() << endl;
						exit(2);
					}
					
					cutOutInfo coi;
					coi.r = r;
					coi.pathNr = toSort[j].second;
					coi.imageCutOut = imageCutOut;
					toAdd.push_back(coi);
					
				}
			}

			for ( unsigned int j = 0 ; j < toAdd.size() ; j++ ) {
				cutOutInfo coi = toAdd[j];
				img::copy( coi.imageCutOut, im, coi.r.xmin, coi.r.ymin );
				img::drawBox( im, coi.r.xmin, coi.r.ymin, coi.r.xmax, coi.r.ymax, colors[ (coi.pathNr % colors.size() )]);
				img::drawBox( im, coi.r.xmin+1, coi.r.ymin+1, coi.r.xmax-1, coi.r.ymax-1, colors[ (coi.pathNr % colors.size() )]);
				img::drawBox( im, coi.r.xmin+2, coi.r.ymin+2, coi.r.xmax-2, coi.r.ymax-2, colors[ (coi.pathNr % colors.size() )]);
			}

			imageVec.push_back( im );

		}
		
		++progress;
		img::ImageRGB bigIm( imageVec[0].width() * 2 , imageVec[0].height() * 2 );
		if ( po.num_of_cams == 3 )
		{
			img::copy( imageVec[0] , bigIm , 0 , 0 );
			img::copy( imageVec[1] , bigIm , imageVec[1].width() , 0 );
			img::copy( imageVec[2] , bigIm , imageVec[2].width() / 2 , imageVec[2].height() );
		} 
		else if ( po.num_of_cams == 4 )
		{
			img::copy( imageVec[0] , bigIm , 0 , 0 );
			img::copy( imageVec[1] , bigIm , imageVec[1].width() , 0 );
			img::copy( imageVec[2] , bigIm , 0 , imageVec[1].height() );
			img::copy( imageVec[3] , bigIm , imageVec[1].width() , imageVec[1].height() );
		}
		img::resample( bigIm , int(imageVec[0].width()*1.5) , int(imageVec[0].height()*1.5) , true );

		string filename = po.output_path + frameNumber + ".jpg";
		img::io::write( filename , bigIm );
	}
}

int main(int argc, char* argv[])
{
	time_t start,end;
	time (&start);
	if( argc != 2 )
	{
		cerr << "[main] Invalid number of arguments. Please (only) specify config file." << endl;
		exit(-1);
	}

	PO po( argv[1] );
	vector<vector<location>> complete_paths;
	int last_time_step = -1;
	int batch_nr = 0;

	for( int begin_frame_counter = po.begin_frame ; begin_frame_counter < po.begin_frame + po.num_of_frames ; 
		begin_frame_counter += po.batch_size ) 
	{
		batch_nr++;
		cout << "Batch nr. " << batch_nr << " of " << po.num_of_batches << "." << endl;

		vector<location> sources;
		getSourceNodes( complete_paths, last_time_step, sources );
		
		Batch b( po, begin_frame_counter, sources );

		if( po.use_appearance_shitrit )
		{
			b.shitrit();
		}
		
		if( batch_nr != po.num_of_batches )
		{
			po.probabilities.pop_back();
		}

		combinePaths( po, b.paths_shitrit, complete_paths, last_time_step );
		last_time_step += po.batch_size;
	}

	if( ! boost::filesystem::exists( po.output_path ) ) 
	{
		boost::filesystem::create_directory( po.output_path );
	}
	
	time (&end);
	double diff = difftime (end,start);
	cout << "Time taken: " << diff << endl;

	vector<vector<int>> converted_paths;
	savePaths( po, complete_paths , converted_paths );

	if ( po.b_save_images )
	{
		saveImages( po, converted_paths );
	}

	return 0;
}

